from call_llm import _run_all_configs_once
from langgraph.graph import StateGraph, START, END
from typing_extensions import TypedDict
from typing import List, Dict, Optional
from pathlib import Path
import sys
import os
import yaml
import asyncio
from fastapi import UploadFile
import re

# 添加项目根目录到路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(current_dir)
sys.path.insert(0, project_root)

# 加载配置文件
CONFIG = {}
try:
    config_path = os.path.join(project_root, "config.yaml")
    if os.path.exists(config_path):
        with open(config_path, "r", encoding="utf-8") as f:
            CONFIG = yaml.safe_load(f) or {}
            if not isinstance(CONFIG, dict):
                CONFIG = {}
except Exception as e:
    print(f"加载 config.yaml 失败：{e}")


class FileJudgementState(TypedDict):
    """文件名判断状态定义"""
    files: Optional[List[UploadFile]]  # 文件列表（用于获取文件名）
    input_file_contents: Optional[List[str]]  # 传入的文件内容列表（不解析，直接使用）
    file_contents: Optional[List[str]]  # 过滤后的文件内容列表
    file_types: Optional[List[str]]  # 文件类型列表（履历/线索）
    file_names: Optional[List[str]]  # 文件名列表（与file_contents和file_types对应）
    result_dict: Optional[Dict[str, List[Dict[str, str]]]]  # 按类型分组的文件内容字典，每个元素包含content和filename
    llm_flag: Optional[str]  # 是否启用LLM判断
    llm_summarize_flag: Optional[str]  # 是否启用LLM总结摘要（true/false）
    length: Optional[int]  # 总结摘要的长度（默认1000）
    delay: Optional[int]  # 并发延迟时间（毫秒，默认100ms）


# async def read_file_content(file: UploadFile) -> str:
#     """
#     读取文件内容，返回文本字符串
#     支持 doc, docx, excel 文件（使用 langchain 解析）
#     其他文件类型尝试 UTF-8 解码，失败则抛出异常
#     """
#     filename = file.filename or ""
#     file_extension = os.path.splitext(filename)[1].lower()
#
#     # 如果是支持的文档类型，使用 langchain 解析
#     if file_extension in [".doc", ".docx", ".xls", ".xlsx"]:
#         try:
#             # parse_document_content 函数内部会处理文件指针重置
#             return await parse_document_content(file)
#         except Exception as e:
#             print(f"Warning: Failed to parse document with langchain: {e}")
#             # 如果解析失败，抛出异常
#             raise ValueError(f"Failed to parse document file {filename}: {e}")
#
#     # 对于其他文件类型，尝试 UTF-8 解码
#     content = await file.read()
#     try:
#         return content.decode('utf-8')
#     except UnicodeDecodeError:
#         raise ValueError(f"File {filename} is not a valid UTF-8 text file and cannot be decoded")


async def judge_file_by_llm(content: str, state: FileJudgementState, idx: int = 0) -> str:
    """
    使用LLM判断文件类型（履历或线索）
    
    Args:
        content: 文件内容
        state: 状态对象
        idx: 文件索引（用于计算延迟）
    
    Returns:
        文件类型（履历/线索/其他）
    """
    delay = state.get("delay", 100)  # 延迟时间（毫秒）
    # 延迟按递增顺序：0, 1*delay, 2*delay, 3*delay...
    delay_ms = idx * delay
    # 延迟执行
    if delay_ms > 0:
        await asyncio.sleep(delay_ms / 1000.0)  # 转换为秒
    
    system_prompt = "你是一个文件分类助手。请根据文件内容判断这是'履历'还是'线索'，只返回'履历'或'线索'两个字。#思考模式 no_think"
    user_prompt = f"请判断以下文件内容是'履历'还是'线索'：\n\n{content[:38000]}"  # 限制长度
    
    results = await _run_all_configs_once(
        name="qwen3-14b",
        system_prompt=system_prompt,
        user_prompt=user_prompt,
        concurrent=1
    )

    if results and len(results) > 0:
        result = results[0] if isinstance(results[0], str) else str(results[0])
        # 去掉<think>\n\n</think>\n\n
        result = re.sub(r'<think>.*?</think>', '', result, flags=re.DOTALL).strip()
        # 提取"履历"或"线索"
        if "履历" in result:
            return "履历"
        elif "线索" in result:
            return "线索"
    
    return "其他"


async def file_name_judgement_node(state: FileJudgementState) -> FileJudgementState:
    """
    节点1：文件名判断和过滤
    如果 llm_flag=true：所有文件都用 LLM 判断（忽略文件名）
    如果 llm_flag=false：
        - 文件名包含"履历" -> 直接判断为"履历"
        - 文件名包含"线索" -> 直接判断为"线索"
        - 其他情况 -> 用 LLM 判断
    """
    files = state.get("files", [])
    input_file_contents = state.get("input_file_contents", [])
    llm_flag = state.get("llm_flag", "false")
    delay = state.get("delay", 100)  # 延迟时间（毫秒）
    result_list = []
    type_list = []
    name_list = []  # 文件名列表
    
    # 确保文件列表和内容列表长度一致
    if len(files) != len(input_file_contents):
        raise ValueError(f"文件数量({len(files)})与文件内容数量({len(input_file_contents)})不匹配")
    
    # 创建并发任务列表和结果占位符
    tasks = []
    task_info = []  # 记录每个任务对应的文件信息（idx, filename, content）
    file_results = [None] * len(files)  # 预分配结果列表，保持原始顺序
    
    # 第一遍遍历：确定哪些需要LLM判断，哪些可以直接判断
    for idx, (file, content) in enumerate(zip(files, input_file_contents)):
        filename = file.filename or ""
        
        if llm_flag.lower() == "true":
            # llm_flag=true：所有文件都用 LLM 判断（并发）
            # 使用 tasks 列表的索引作为延迟计算的索引，延迟按递增顺序：0, 1*delay, 2*delay...
            task_idx = len(tasks)
            task = judge_file_by_llm(content, state, task_idx)
            tasks.append(task)
            task_info.append((idx, filename, content))
        else:
            # llm_flag=false：先根据文件名判断，如果文件名中没有"履历"或"线索"，再用 LLM 判断
            if "履历" in filename:
                file_type = "履历"
                file_results[idx] = (content, file_type, filename)  # 保存文件名
            elif "线索" in filename:
                file_type = "线索"
                file_results[idx] = (content, file_type, filename)  # 保存文件名
            else:
                # 文件名中不包含"履历"也不包含"线索"，使用 LLM 判断（并发）
                # 使用 tasks 列表的索引作为延迟计算的索引，延迟按递增顺序：0, 1*delay, 2*delay...
                task_idx = len(tasks)
                task = judge_file_by_llm(content, state, task_idx)
                tasks.append(task)
                task_info.append((idx, filename, content))
    
    # 并发执行所有LLM判断任务
    if tasks:
        llm_results = await asyncio.gather(*tasks)
        
        # 处理LLM判断结果，填充到对应位置
        for (idx, filename, content), file_type in zip(task_info, llm_results):
            # 如果是履历或线索，保存结果
            if file_type in ["履历", "线索"]:
                file_results[idx] = (content, file_type, filename)  # 保存文件名
    
    # 按原始顺序组装结果  [('我是娃哈哈', '履历', 'filename.txt')] 多个元组
    for result in file_results:
        if result is not None:
            content, file_type, filename = result
            result_list.append(content)
            type_list.append(file_type)
            name_list.append(filename)
    
    return {
        "file_contents": result_list,
        "file_types": type_list,
        "file_names": name_list
    }


async def llm_summarize_node(state: FileJudgementState) -> FileJudgementState:
    """
    节点2：LLM总结摘要节点
    对履历和线索文件分别进行总结摘要，串行遍历执行（不使用并发）
    """
    result_dict = state.get("result_dict", {})
    length = state.get("length", 1000)
    
    summarized_dict = {
        "履历": [],
        "线索": []
    }
    
    # ========== 串行执行方式（当前使用） ==========
    # 串行执行：遍历履历文件，逐个执行总结摘要（不使用并发，不需要延迟）
    if "履历" in result_dict and result_dict["履历"]:
        for idx, item in enumerate(result_dict["履历"]):
            # item 是 {"content": "...", "filename": "..."}
            content = item.get("content", "")
            filename = item.get("filename", "")
            # 直接调用总结摘要函数，不使用延迟，串行执行
            summarized_content = await summarize_content(content, "履历", length)
            # 保存结果，包含文件名
            summarized_dict["履历"].append({
                "content": summarized_content,
                "filename": filename
            })
    
    # 串行执行：遍历线索文件，逐个执行总结摘要（不使用并发，不需要延迟）
    if "线索" in result_dict and result_dict["线索"]:
        for idx, item in enumerate(result_dict["线索"]):
            # item 是 {"content": "...", "filename": "..."}
            content = item.get("content", "")
            filename = item.get("filename", "")
            # 直接调用总结摘要函数，不使用延迟，串行执行
            summarized_content = await summarize_content(content, "线索", length)
            # 保存结果，包含文件名
            summarized_dict["线索"].append({
                "content": summarized_content,
                "filename": filename
            })
    
    # ========== 并发执行方式（已注释，保留参考） ==========
    # delay = state.get("delay", 100)  # 延迟时间（毫秒）
    # 
    # # 创建并发任务列表（记录类型和索引）
    # tasks = []
    # task_info = []  # 记录每个任务对应的类型和索引
    # 
    # # 处理履历文件（并发）
    # if "履历" in result_dict and result_dict["履历"]:
    #     for idx, content in enumerate(result_dict["履历"]):
    #         # 使用 tasks 列表的索引作为延迟计算的索引，延迟按递增顺序：0, 1*delay, 2*delay...
    #         task_idx = len(tasks)
    #         task = summarize_with_delay(content, "履历", state, task_idx)
    #         tasks.append(task)
    #         task_info.append(("履历", idx))
    # 
    # # 处理线索文件（并发）
    # if "线索" in result_dict and result_dict["线索"]:
    #     for idx, content in enumerate(result_dict["线索"]):
    #         # 使用 tasks 列表的索引作为延迟计算的索引，延迟按递增顺序：0, 1*delay, 2*delay...
    #         task_idx = len(tasks)
    #         task = summarize_with_delay(content, "线索", state, task_idx)
    #         tasks.append(task)
    #         task_info.append(("线索", idx))
    # 
    # # 并发执行所有任务
    # if tasks:
    #     results = await asyncio.gather(*tasks)
    #     # 按类型分组结果（保持原始顺序）
    #     for (file_type, idx), result in zip(task_info, results):
    #         # 确保列表有足够的长度
    #         if len(summarized_dict[file_type]) <= idx:
    #             # 扩展列表到所需长度
    #             summarized_dict[file_type].extend([""] * (idx + 1 - len(summarized_dict[file_type])))
    #         summarized_dict[file_type][idx] = result
    
    return {"result_dict": summarized_dict}


async def summarize_with_delay(content: str, file_type: str, state: FileJudgementState, idx: int = 0) -> str:
    """
    带延迟的总结摘要函数
    
    Args:
        content: 原始内容
        file_type: 文件类型（履历/线索）
        state: 状态对象
        idx: 文件索引（用于计算延迟）
    
    Returns:
        总结摘要后的内容
    """
    delay = state.get("delay", 100)  # 延迟时间（毫秒）
    length = state.get("length", 1000)  # 摘要长度
    delay_ms = idx * delay
    # 延迟执行
    if delay_ms > 0:
        await asyncio.sleep(delay_ms / 1000.0)  # 转换为秒
    
    # 调用总结摘要
    return await summarize_content(content, file_type, length)


async def summarize_content(content: str, file_type: str, length: int) -> str:
    """
    使用LLM对内容进行总结摘要
    
    Args:
        content: 原始内容
        file_type: 文件类型（履历/线索）
        length: 摘要长度
    
    Returns:
        总结摘要后的内容
    """
    system_prompt = f"""你是一个专业的文档摘要助手。请对以下{file_type}文件内容进行总结摘要，摘要长度控制在{length}字左右。
                    
                    # 要求如下：
                    1. 保留关键信息，去除冗余内容
                    2. 保持原文的逻辑结构和重要信息
                    3. 使用简洁明了的语言表达
                    4. 不要扩展内容，只保留原文的精华部分
        
                    #思考模式 
                    no_think"""

    user_prompt = f"请对以下{file_type}文件内容进行总结摘要：\n\n{content[:38000]}"  # 限制输入长度
    
    results = await _run_all_configs_once(
        name="qwen3-14b",
        system_prompt=system_prompt,
        user_prompt=user_prompt,
        concurrent=1
    )
    
    if results and len(results) > 0:
        result = results[0] if isinstance(results[0], str) else str(results[0])
        # 去掉<think>\n\n</think>\n\n
        result = re.sub(r'<think>.*?</think>', '', result, flags=re.DOTALL).strip()
        return result
    
    # 如果LLM调用失败，返回原始内容
    return content


def end_node(state: FileJudgementState) -> FileJudgementState:
    """
    结束节点：将文件内容按类型分组，包含文件名
    """
    file_contents = state.get("file_contents", [])
    file_types = state.get("file_types", [])
    file_names = state.get("file_names", [])
    
    # 按类型分组，每个元素包含content和filename
    result_dict = {
        "履历": [],
        "线索": []
    }
    # zip 会将三个列表按索引位置配对
    for content, file_type, filename in zip(file_contents, file_types, file_names):
        if file_type in result_dict:
            result_dict[file_type].append({
                "content": content,
                "filename": filename
            })
    
    return {"result_dict": result_dict}


# 构建LangGraph
# 注意：节点必须预先定义，不能根据运行时 state 动态添加
# 是否执行 llm_summarize_node 通过条件边 should_summarize 来控制
graph_builder = StateGraph(FileJudgementState)
graph_builder.add_node("file_name_judgement", file_name_judgement_node)
graph_builder.add_node("llm_summarize", llm_summarize_node)  # 必须添加节点，通过条件边控制是否执行
graph_builder.add_node("end", end_node)

# 定义边
graph_builder.add_edge(START, "file_name_judgement")
graph_builder.add_edge("file_name_judgement", "end")

# 根据 llm_summarize_flag 决定是否执行 llm_summarize 节点
def should_summarize(state: FileJudgementState) -> str:
    """判断是否应该执行总结摘要"""
    llm_summarize_flag = state.get("llm_summarize_flag", "false")
    
    # 调试日志：打印 llm_summarize_flag 的值和类型
    print(f"[DEBUG] should_summarize: llm_summarize_flag = {llm_summarize_flag}, type = {type(llm_summarize_flag)}")
    
    # 处理各种类型的值：字符串、布尔值、None等
    if llm_summarize_flag is None:
        print(f"[DEBUG] should_summarize: llm_summarize_flag is None, returning END")
        return "END"
    
    # 如果是布尔值，直接判断
    if isinstance(llm_summarize_flag, bool):
        if llm_summarize_flag:
            print(f"[DEBUG] should_summarize: llm_summarize_flag is True (bool), returning llm_summarize")
            return "llm_summarize"
        else:
            print(f"[DEBUG] should_summarize: llm_summarize_flag is False (bool), returning END")
            return "END"
    
    # 如果是字符串，转换为小写后判断
    if isinstance(llm_summarize_flag, str):
        flag_lower = llm_summarize_flag.lower()
        if flag_lower == "true":
            print(f"[DEBUG] should_summarize: llm_summarize_flag is 'true' (str), returning llm_summarize")
            return "llm_summarize"
        else:
            print(f"[DEBUG] should_summarize: llm_summarize_flag is '{llm_summarize_flag}' (str), returning END")
            return "END"
    
    # 其他情况，默认不执行
    print(f"[DEBUG] should_summarize: llm_summarize_flag is other type, returning END")
    return "END"

graph_builder.add_conditional_edges(
    "end",
    should_summarize,
    {
        "llm_summarize": "llm_summarize",
        "END": END
    }
)
graph_builder.add_edge("llm_summarize", END)

# 编译图
graph = graph_builder.compile()


async def file_name_judgement(
    filenames: List[str], 
    file_contents: List[str], 
    llm_flag: str = "false",
    llm_summarize_flag: str = "false",
    length: int = 1000,
    delay: int = 100
) -> List[Dict[str, str]]:
    """
    根据文件名过滤文件，返回扁平化的文件类型和内容列表（使用LangGraph）
    
    Args:
        filenames: 文件名列表（字符串数组）
        file_contents: 文件内容列表（已解析好的内容，不在这里解析）
        llm_flag: 是否启用LLM判断
        llm_summarize_flag: 是否启用LLM总结摘要（true/false字符串）
        length: 总结摘要的长度（默认1000）
        delay: 并发延迟时间（毫秒，默认100ms）
    
    Returns:
        返回格式：[{"type": "履历", "content": "...", "filename": "..."}, ...]
        方便提取列表：
        - list1 = [item["type"] for item in result]  # 文件类型列表
        - list2 = [item["content"] for item in result]  # 内容列表
        - list3 = [item["filename"] for item in result]  # 文件名列表
    """
    # 创建模拟的 UploadFile 对象列表（只用于获取文件名）
    from io import BytesIO
    
    files = []
    for filename in filenames:
        # 创建一个空的 UploadFile 对象，只包含文件名
        file_obj = UploadFile(
            filename=filename,
            file=BytesIO(b""),  # 空文件内容
            headers={"content-type": "application/octet-stream"}
        )
        files.append(file_obj)
    
    initial_state = {
        "files": files,
        "input_file_contents": file_contents,
        "llm_flag": llm_flag,
        "llm_summarize_flag": llm_summarize_flag,
        "length": length,
        "delay": delay,
        "file_contents": None,
        "file_types": None,
        "file_names": None,
        "result_dict": None
    }
    
    # 运行LangGraph（使用异步调用）
    final_state = await graph.ainvoke(initial_state)
    
    # 获取结果字典
    result_dict = final_state.get("result_dict", {"履历": [], "线索": []})
    
    # 转换为扁平化列表格式：[{"type": "履历", "content": "...", "filename": "..."}, ...]
    result_list = []
    for file_type in ["履历", "线索"]:
        if file_type in result_dict and result_dict[file_type]:
            for item in result_dict[file_type]:
                # item 是 {"content": "...", "filename": "..."}
                result_list.append({
                    "type": file_type,
                    "content": item.get("content", ""),
                    "filename": item.get("filename", "")
                })
    
    return result_list